using System;
using System.Collections.Generic;
using System.Text;
using DotNetMatrix;

namespace Blimpbots.Recognition
{
    public class ProjectiveCamera
    {
        public CameraProperties Properties;
        protected GeneralMatrix _k;
        protected GeneralMatrix _hInv; // homogeneous transformation inverted

        public ProjectiveCamera(CameraProperties properties)
        {
            if(properties.intrinsicMatrix != null)
                _k = new GeneralMatrix(properties.intrinsicMatrix);

            this.Properties = properties;
            UpdateExtrinsic();
        }
        
        public Point3D EstimateCoordinate(Point2D imagePoint, double depth, out GeneralMatrix ellipseRotation)
        {
            /* needed intrinsic params */
            double
                fx = _k.Array[0][0],
                fy = _k.Array[1][1],
                cx = _k.Array[0][2],
                cy = _k.Array[1][2];

            /* the line out of the camera center is just
             * lambda * _k.Inverse() * [ x y 1 ]'
             * with lambda as a parameter along the line
             * which ends up giving the form below */

            double a = (imagePoint.x - cx) / fx;
            double b = (imagePoint.y - cy) / fy;
            double lambda = Math.Sqrt(depth * depth / (a * a + b * b + 1));

            GeneralMatrix rotX = RotationMatrix(Math.Atan2(a, 1), RotationType.PitchY);
            GeneralMatrix rotY = RotationMatrix(Math.Atan2(b, 1), RotationType.RollX);

            /* Rotation from world-space to error-ellipse */
            ellipseRotation = rotX * rotY * _hInv.GetMatrix(0, 2, 0, 2);
 
            GeneralMatrix pt3d = new Point3D(lambda * a, lambda * b, lambda).ToMatrix();
            pt3d = _hInv.Multiply(pt3d);

            return new Point3D(pt3d.Array[0][0],
                               pt3d.Array[1][0],
                               pt3d.Array[2][0]);
        }

        private enum RotationType { RollX, PitchY, YawZ }
        private GeneralMatrix RotationMatrix(double theta, RotationType type)
        {
            double sin = Math.Sin(theta);
            double cos = Math.Cos(theta);

            switch (type)
            {
                case RotationType.RollX:
                    return new GeneralMatrix(new double[] { 1, 0, 0, 0, cos, sin, 0, -sin, cos }, 3);
                case RotationType.PitchY:
                    return new GeneralMatrix(new double[] { cos, 0, -sin, 0, 1, 0, sin, 0, cos }, 3);
                default: // PitchZ
                    return new GeneralMatrix(new double[] { cos, sin, 0, -sin, cos, 0,  0, 0 ,1}, 3);
            }
        }

        public double EstimateDepth(Point2D imagePoint, double z)
        {
            /* needed intrinsic params */
            double
                fx = _k.Array[0][0],
                fy = _k.Array[1][1],
                cx = _k.Array[0][2],
                cy = _k.Array[1][2];

            double a = (imagePoint.x - cx) / fx;
            double b = (imagePoint.y - cy) / fy;
            double lambda = Math.Sqrt(z * z / (a * a + b * b + 1));

            return lambda;
        }


        /// <summary>
        /// Returns the world-reference-frame coordinate of a camera-frame point
        /// </summary>
        /// <param name="camSpacePt">Camera-frame point.</param>
        /// <returns>World-frame point.</returns>
        public Point3D ConvertCoordinate(Point3D camFramePt)
        {
            return Point3D.FromMatrix(
                _hInv.Multiply(camFramePt.ToMatrix()));
        }

        /// <summary>
        /// Undistorts image points using the Oulu University
        /// method.
        /// </summary>
        /// <param name="pt">Distorted image points.</param>
        /// <returns>Undistorted image points.</returns>
        public Point2D Undistort(Point2D pt)
        {
            if (Properties.distortion == null) return pt;

            double ifx = 1.0 / Properties.FX;
            double ify = 1.0 / Properties.FY;

            double[] k = Properties.distortion;

            double x = (pt.x - Properties.CX) * ifx,
                   y = (pt.y - Properties.CY) * ify,
                   x0 = x, y0 = y;

            double k1 = Properties.distortion[0];
            double k2 = Properties.distortion[1];
            double p1 = Properties.distortion[2];
            double p2 = Properties.distortion[3];

            for (int j = 0; j < 20; j++)
            {
                double r2 = x * x + y * y;
                double k_radial = 1 + k1 * r2 + k2 * (r2 * r2);
                double delta_x = 2 * p1 * x * y + p2 * (r2 + 2 * x * x);
                double delta_y = p1 * (r2 + 2 * y * y) + 2 * p2 * x * y;
                x = (x0 - delta_x) / k_radial;
                y = (y0 - delta_y) / k_radial;
            }

            double u = Properties.FX * x + Properties.CX;
            double v = Properties.FY * y + Properties.CY;
            return new Point2D(u, v);

        }

        public Point2D Distort(Point2D pt)
        {
            // As written in Intel OpenCV Undistort2

            double u = pt.x;
            double v = pt.y;

            double u0 = Properties.CX;
            double v0 = Properties.CY;
            double fx = Properties.FX;
            double fy = Properties.FY;

            double k1 = Properties.distortion[0];
            double k2 = Properties.distortion[1];
            double p1 = Properties.distortion[2];
            double p2 = Properties.distortion[3];

            // Y
            double y = (v - v0) * (1.0 / fy);
            double y2 = y * y;
            double ky = 1 + (k1 + k2 * y2) * y2;
            double k2y = 2 * k2 * y2;
            double _2p1y = 2 * p1 * y;
            double _3p1y2 = 3 * p1 * y2;
            double p2y2 = p2 * y2;

            // X
            double x = (u - u0) * (1.0 / fx);
            double x2 = x * x;
            double kx = (k1 + k2 * x2) * x2;
            double d = kx + ky + k2y * x2;

            double _u = fx * (x * (d + _2p1y) + p2y2 + (3 * p2) * x2) + u0;
            double _v = fy * (y * (d + (2 * p2) * x) + _3p1y2 + p1 * x2) + v0;

            return new Point2D(_u, _v);

        }

        public void UpdateExtrinsic()
        {
            /* Can't build transform without knowing this */
            if (Properties.rotationMatrix == null ||
                Properties.translationVector == null)
                return;

            #region Construct homogeneous Camera --> World transformation

            GeneralMatrix r = new GeneralMatrix(Properties.rotationMatrix);
            GeneralMatrix t = new GeneralMatrix(Properties.translationVector, 3);

            GeneralMatrix H = new GeneralMatrix(4, 4);
            H.SetMatrix(0, 2, 0, 2, r);
            H.SetMatrix(0, 2, 3, 3, t);
            H.Array[3][3] = 1.0;

            _hInv = H.Inverse();

            #endregion
        }

    }
}
